{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1ae38945-39dd-45dc-ad4f-da7a4404241f",
   "metadata": {},
   "source": [
    "<table style=\"width:100%\">\n",
    "<tr>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<font size=\"2\">\n",
    "Supplementary code for the <a href=\"http://mng.bz/orYv\">Build a Large Language Model From Scratch</a> book by <a href=\"https://sebastianraschka.com\">Sebastian Raschka</a><br>\n",
    "<br>Code repository: <a href=\"https://github.com/rasbt/LLMs-from-scratch\">https://github.com/rasbt/LLMs-from-scratch</a>\n",
    "<br>汉化的库: <a href=\"https://github.com/GoatCsu/CN-LLMs-from-scratch.git\">https://github.com/GoatCsu/CN-LLMs-from-scratch.git</a>\n",
    "</font>\n",
    "</td>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<a href=\"http://mng.bz/orYv\"><img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/cover-small.webp\" width=\"100px\"></a>\n",
    "</td>\n",
    "</tr>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bfa70ec-5c4c-40e8-b923-16f8167e3181",
   "metadata": {},
   "source": [
    "# 第三章: Attention"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c29bcbe8-a034-43a2-b557-997b03c9882d",
   "metadata": {},
   "source": [
    "本章所需要的包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e58f33e8-5dc9-4dd5-ab84-5a011fa11d92",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch version: 2.6.0+cu126\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "\n",
    "print(\"torch version:\", version(\"torch\"))\n",
    "#导入并确认库"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2a4474d-7c68-4846-8702-37906cf08197",
   "metadata": {},
   "source": [
    "- LLM的核心:Attention\n",
    "- 译者:可以直接看论文呀!\n",
    "- [Attention is all you need](https://arxiv.org/abs/1706.03762)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02a11208-d9d3-44b1-8e0d-0c8414110b93",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/01.webp?123\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50e020fd-9690-4343-80df-da96678bef5e",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/02.webp\" width=\"600px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecc4dcee-34ea-4c05-9085-2f8887f70363",
   "metadata": {},
   "source": [
    "## 3.1 长序列的建模"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a55aa49c-36c2-48da-b1d9-70f416e46a6a",
   "metadata": {},
   "source": [
    "- 没有代码\n",
    "- 逐字翻译文本通常不可行，因为源语言和目标语言在语法结构上存在差异："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55c0c433-aa4b-491e-848a-54905ebb05ad",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/03.webp\" width=\"400px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db03c48a-3429-48ea-9d4a-2e53b0e516b1",
   "metadata": {},
   "source": [
    "- 在Transformer模型出现之前，机器翻译任务主要依赖于编码器(encoder)-解码器(decoder)架构的循环神经网络（RNNs）。\n",
    "- 在这种架构中，编码器逐词处理源语言序列，并通过隐藏状态（神经网络中的中间层）生成输入序列的表示："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03d8df2c-c1c2-4df0-9977-ade9713088b2",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/04.webp\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3602c585-b87a-41c7-a324-c5e8298849df",
   "metadata": {},
   "source": [
    "## 3.2 注意力机制高效捕获数据关系"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6fde64c-6034-421d-81d9-8244932086ea",
   "metadata": {},
   "source": [
    "- 本节不涉及代码。\n",
    "- 借助注意力机制，文本生成解码器能够选择性地关注所有输入token，从而在生成特定输出token时，动态分配不同输入token的重要性系数"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc4f6293-8ab5-4aeb-a04c-50ee158485b1",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/05.webp\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8044be1f-e6a2-4a1f-a6dd-e325d3bad05e",
   "metadata": {},
   "source": [
    "- Transformer中的自注意力机制是一种关键技术，它通过让序列中的每个位置与其他所有位置交互并计算相关性，从而增强输入表示的上下文信息。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6565dc9f-b1be-4c78-b503-42ccc743296c",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/06.webp\" width=\"300px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5efe05ff-b441-408e-8d66-cde4eb3397e3",
   "metadata": {},
   "source": [
    "## 3.3 自注意力关注的不同部分"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d9af516-7c37-4400-ab53-34936d5495a9",
   "metadata": {},
   "source": [
    "### 3.3.1 无可变参数的自注意力模型"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d269e9f1-df11-4644-b575-df338cf46cdf",
   "metadata": {},
   "source": [
    "- 本节介绍了一种高度简化的自注意力变体，不包含任何可训练的权重。\n",
    "- 该变体仅用于说明目的，并非Transformer中实际使用的注意力机制。\n",
    "- 下一节（3.3.2节）将扩展此简易模型，实现真正的自注意力机制。\n",
    "- 假设给定一个输入序列 $x^{(1)}$ 到 $x^{(T)}$：\n",
    "  - 输入是一个文本（例如，一句已被处理为token嵌入的句子，如“Your journey starts with one step”），具体处理方法在第2章中已有描述。\n",
    "  - 例如，$x^{(1)}$ 是表示单词“Your”的d维向量，以此类推。\n",
    "\n",
    "- **目标：** 为输入序列中的每个元素 $x^{(i)}$（从 $x^{(1)}$ 到 $x^{(T)}$）计算上下文向量 $z^{(i)}$（$z$ 和 $x$ 的维度相同）。\n",
    "    - 上下文向量 $z^{(i)}$ 是对输入 $x^{(1)}$ 到 $x^{(T)}$ 的加权求和。\n",
    "    - 上下文向量是针对特定输入的“上下文”相关表示。\n",
    "      - 以第二个输入 $x^{(2)}$ 为例，说明具体计算过程。\n",
    "      - 第二个上下文向量 $z^{(2)}$ 是对所有输入 $x^{(1)}$ 到 $x^{(T)}$ 的加权求和，权重由相对于 $x^{(2)}$ 的注意力权重决定。\n",
    "      - 注意力权重决定了每个输入元素对 $z^{(2)}$ 的贡献程度。\n",
    "      - 简而言之，$z^{(2)}$ 是 $x^{(2)}$ 的增强版本，融合了与当前任务相关的所有其他输入元素的信息。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcc7c7a2-b6ab-478f-ae37-faa8eaa8049a",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/07.webp\" width=\"400px\">\n",
    "\n",
    "- （请注意，此图中的数字已截断至小数点后一位，以减少视觉干扰；其他图表中的数值也可能经过类似处理。）"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff856c58-8382-44c7-827f-798040e6e697",
   "metadata": {},
   "source": [
    "- 按照惯例，未归一化的注意力值称为 **“注意力得分”**，而归一化后总和为1的注意力得分称为 **“注意力权重”**。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01b10344-128d-462a-823f-2178dff5fd58",
   "metadata": {},
   "source": [
    "- 下方代码逐步演示了上图的操作过程\n",
    "\n",
    "<br>\n",
    "\n",
    "- **步骤 1：** 计算未归一化的注意力得分 $\\omega$\n",
    "- 假设使用第二个输入token作为查询，即 $q^{(2)} = x^{(2)}$，通过点积计算未归一化的注意力得分：\n",
    "    - $\\omega_{21} = x^{(1)} \\cdot q^{(2)\\top}$\n",
    "    - $\\omega_{22} = x^{(2)} \\cdot q^{(2)\\top}$\n",
    "    - $\\omega_{23} = x^{(3)} \\cdot q^{(2)\\top}$\n",
    "    - ...\n",
    "    - $\\omega_{2T} = x^{(T)} \\cdot q^{(2)\\top}$\n",
    "- 其中，$\\omega$ 是希腊字母“欧米伽”，表示未归一化的注意力得分。\n",
    "    - 在 $\\omega_{21}$ 中，下标“21”表示以第2个元素为查询，与第1个元素计算得分。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35e55f7a-f2d0-4f24-858b-228e4fe88fb3",
   "metadata": {},
   "source": [
    "- 假设我们有以下输入句子，该句子已根据第3章的描述嵌入到3维向量中（此处我们使用了一个非常小的嵌入维度进行说明，以便内容可以显示在页面上）： "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "22b9556a-aaf8-4ab4-a5b4-973372b0b2c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "inputs = torch.tensor(\n",
    "  [[0.43, 0.15, 0.89], # Your     (x^1)\n",
    "   [0.55, 0.87, 0.66], # journey  (x^2)\n",
    "   [0.57, 0.85, 0.64], # starts   (x^3)\n",
    "   [0.22, 0.58, 0.33], # with     (x^4)\n",
    "   [0.77, 0.25, 0.10], # one      (x^5)\n",
    "   [0.05, 0.80, 0.55]] # step     (x^6)\n",
    ")\n",
    "#对于一句话中的每个单词定义了一个三维的向量"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "299baef3-b1a8-49ba-bad4-f62c8a416d83",
   "metadata": {},
   "source": [
    "- （在本书中，我们遵循机器学习和深度学习的常见惯例：训练样本以行表示，特征值以列表示；对于上述张量，每一行表示一个词，每一列表示一个嵌入维度。）\n",
    "\n",
    "- 本节的主要目标是演示如何以第二个输入序列 $x^{(2)}$ 作为查询，计算其上下文向量 $z^{(2)}$。\n",
    "\n",
    "- 图中展示了该过程的第一步，即通过点积操作计算 $x^{(2)}$ 与所有其他输入元素之间的注意力得分 $\\omega$。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cb3453a-58fa-42c4-b225-86850bc856f8",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/08.webp\" width=\"400px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77be52fb-82fd-4886-a4c8-f24a9c87af22",
   "metadata": {},
   "source": [
    "- 我们以输入序列中的第2个元素 $x^{(2)}$ 为例，计算其上下文向量 $z^{(2)}$；稍后会将此方法推广至计算所有上下文向量。\n",
    "- 第一步是通过计算查询 $x^{(2)}$ 与所有输入token的点积，得到未归一化的注意力得分："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6fb5b2f8-dd2c-4a6d-94ef-a0e9ad163951",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])\n"
     ]
    }
   ],
   "source": [
    "query = inputs[1]  # 2nd input token is the query\n",
    "\n",
    "attn_scores_2 = torch.empty(inputs.shape[0])\n",
    "#建立一个未初始化的张量来记录注意力得分\n",
    "for i, x_i in enumerate(inputs):\n",
    "    attn_scores_2[i] = torch.dot(x_i, query) \n",
    "    # 相似性度量计算attention分数\n",
    "    # 从公式上看也就是点乘\n",
    "\n",
    "print(attn_scores_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8df09ae0-199f-4b6f-81a0-2f70546684b8",
   "metadata": {},
   "source": [
    "- 补充说明：点积本质上是逐元素相乘并将所得积相加的一种简写表示："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9842f39b-1654-410e-88bf-d1b899bf0241",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.9544)\n",
      "tensor(0.9544)\n"
     ]
    }
   ],
   "source": [
    "res = 0.\n",
    "\n",
    "for idx, element in enumerate(inputs[0]):\n",
    "    res += inputs[0][idx] * query[idx]\n",
    "    #累加Key*Query的成绩\n",
    "\n",
    "print(res)\n",
    "print(torch.dot(inputs[0], query))\n",
    "#相当于解释了一遍点成的内部原理"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d444d76-e19e-4e9a-a268-f315d966609b",
   "metadata": {},
   "source": [
    "- **步骤 2：** 将未归一化的注意力得分（“欧米伽”，$\\omega$）归一化，使其总和为1。\n",
    "- 以下是一种简单的方法，用于将未归一化的注意力得分归一化："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfd965d6-980c-476a-93d8-9efe603b1b3b",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/09.webp\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e3ccc99c-33ce-4f11-b7f2-353cf1cbdaba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attention weights: tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656])\n",
      "Sum: tensor(1.0000)\n"
     ]
    }
   ],
   "source": [
    "attn_weights_2_tmp = attn_scores_2 / attn_scores_2.sum() \n",
    "#归一化,这里是属于加权式的归一化\n",
    "\n",
    "print(\"Attention weights:\", attn_weights_2_tmp)\n",
    "print(\"Sum:\", attn_weights_2_tmp.sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75dc0a57-f53e-41bf-8793-daa77a819431",
   "metadata": {},
   "source": [
    "- 然而，在实践中，使用softmax函数进行归一化更为常见，因为它能够更好地处理极端值，并且在训练过程中具有更理想的梯度特性，因此推荐使用。\n",
    "- 下面是一个简单的softmax函数实现，用于缩放并对向量元素进行归一化，使它们的和为1："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "07b2e58d-a6ed-49f0-a1cd-2463e8d53a20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])\n",
      "Sum: tensor(1.)\n"
     ]
    }
   ],
   "source": [
    "def softmax_naive(x):\n",
    "    return torch.exp(x) / torch.exp(x).sum(dim=0)\n",
    "\n",
    "attn_weights_2_naive = softmax_naive(attn_scores_2)\n",
    "\n",
    "print(\"Attention weights:\", attn_weights_2_naive)\n",
    "print(\"Sum:\", attn_weights_2_naive.sum())\n",
    "##用SoftMax做归一化, 处理好极端值\n",
    "#有合理的梯度数据表现力"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0a1cbbb-4744-41cb-8910-f5c1355555fb",
   "metadata": {},
   "source": [
    "- 上述简单实现可能会因输入值过大或过小而遭遇数值不稳定问题，导致溢出或下溢。\n",
    "- 因此，在实践中，建议使用PyTorch内置的softmax函数，因为它经过高度优化，性能更佳："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2d99cac4-45ea-46b3-b3c1-e000ad16e158",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])\n",
      "Sum: tensor(1.)\n"
     ]
    }
   ],
   "source": [
    "attn_weights_2 = torch.softmax(attn_scores_2, dim=0)\n",
    "\n",
    "print(\"Attention weights:\", attn_weights_2)\n",
    "print(\"Sum:\", attn_weights_2.sum())\n",
    "#用torch优化过的softmax对边缘值也挺友好的"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e43e36c7-90b2-427f-94f6-bb9d31b2ab3f",
   "metadata": {},
   "source": [
    "- **步骤 3**：通过将嵌入的输入标记 $x^{(i)}$ 与注意力权重相乘，并对结果向量求和，计算上下文向量 $z^{(2)}$："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1c9f5ac-8d3d-4847-94e3-fd783b7d4d3d",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/10.webp\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8fcb96f0-14e5-4973-a50e-79ea7c6af99f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.4419, 0.6515, 0.5683])\n"
     ]
    }
   ],
   "source": [
    "query = inputs[1] # 2nd input token is the query\n",
    "\n",
    "context_vec_2 = torch.zeros(query.shape)\n",
    "#创造一个内容的零向量\n",
    "for i,x_i in enumerate(inputs):\n",
    "    context_vec_2 += attn_weights_2[i]*x_i\n",
    "    #把不同内容的向量+起来\n",
    "\n",
    "print(context_vec_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a454262-40eb-430e-9ca4-e43fb8d6cd89",
   "metadata": {},
   "source": [
    "### 3.3.2 计算所有token的attention score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a02bb73-fc19-4c88-b155-8314de5d63a8",
   "metadata": {},
   "source": [
    "#### 将其推广到所有输入序列标记：\n",
    "\n",
    "- 上面，我们为输入2计算了注意力权重和上下文向量（如下图中高亮的行所示）。\n",
    "- 接下来，我们将此计算推广到所有输入序列标记，计算对应的注意力权重和上下文向量。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11c0fb55-394f-42f4-ba07-d01ae5c98ab4",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/11.webp\" width=\"400px\">\n",
    "\n",
    "- （请注意，图中的数字已四舍五入到小数点后两位；每一行的数值应相加为1.0或100%；其他图中的数字也进行了类似处理。）"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b789b990-fb51-4beb-9212-bf58876b5983",
   "metadata": {},
   "source": [
    "- 在自注意力机制中，首先计算注意力得分，随后对这些得分进行归一化，得到总和为1的注意力权重。\n",
    "- 接着，利用这些注意力权重对输入进行加权求和，生成上下文向量。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9bffe4b-56fe-4c37-9762-24bd924b7d3c",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/12.webp\" width=\"400px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa652506-f2c8-473c-a905-85c389c842cc",
   "metadata": {},
   "source": [
    "- 对所有成对元素应用之前的**步骤 1**，计算未归一化的注意力得分矩阵："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "04004be8-07a1-468b-ab33-32e16a551b45",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],\n",
      "        [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],\n",
      "        [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],\n",
      "        [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],\n",
      "        [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],\n",
      "        [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])\n"
     ]
    }
   ],
   "source": [
    "attn_scores = torch.empty(6, 6)\n",
    "#建立个空表来储存相关联程度\n",
    "\n",
    "for i, x_i in enumerate(inputs):\n",
    "    for j, x_j in enumerate(inputs):\n",
    "        attn_scores[i, j] = torch.dot(x_i, x_j)\n",
    "        #一点点计算相关性并输入表格\n",
    "print(attn_scores)\n",
    "#事实上就是实现了两个单词之间的关联度列表输出"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1539187f-1ece-47b7-bc9b-65a97115f1d4",
   "metadata": {},
   "source": [
    "- 如果是矩阵相乘那么更有效率"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2cea69d0-9a47-45da-8d5a-47ceef2df673",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],\n",
      "        [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],\n",
      "        [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],\n",
      "        [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],\n",
      "        [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],\n",
      "        [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])\n"
     ]
    }
   ],
   "source": [
    "attn_scores = inputs @ inputs.T\n",
    "print(attn_scores)\n",
    "#有简单的方法整合方法计算"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02c4bac4-acfd-427f-9b11-c436ac71748d",
   "metadata": {},
   "source": [
    "- 与**第二步**相似, 我们对每一行都要归一化操作:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fa4ef062-de81-47ee-8415-bfe1708c81b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452],\n",
      "        [0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581],\n",
      "        [0.1390, 0.2369, 0.2326, 0.1242, 0.1108, 0.1565],\n",
      "        [0.1435, 0.2074, 0.2046, 0.1462, 0.1263, 0.1720],\n",
      "        [0.1526, 0.1958, 0.1975, 0.1367, 0.1879, 0.1295],\n",
      "        [0.1385, 0.2184, 0.2128, 0.1420, 0.0988, 0.1896]])\n"
     ]
    }
   ],
   "source": [
    "attn_weights = torch.softmax(attn_scores, dim=-1)\n",
    "print(attn_weights)\n",
    "#归一化处理"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fa6d02b-7f15-4eb4-83a7-0b8a819e7a0c",
   "metadata": {},
   "source": [
    "- 一个快速验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "112b492c-fb6f-4e6d-8df5-518ae83363d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Row 2 sum: 1.0\n",
      "All row sums: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000])\n"
     ]
    }
   ],
   "source": [
    "row_2_sum = sum([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])\n",
    "print(\"Row 2 sum:\", row_2_sum)\n",
    "\n",
    "print(\"All row sums:\", attn_weights.sum(dim=-1))\n",
    "#验证一下大家加起来都是1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "138b0b5c-d813-44c7-b373-fde9540ddfd1",
   "metadata": {},
   "source": [
    "- 用**step 3** 计算所有的向量:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ba8eafcf-f7f7-4989-b8dc-61b50c4f81dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.4421, 0.5931, 0.5790],\n",
      "        [0.4419, 0.6515, 0.5683],\n",
      "        [0.4431, 0.6496, 0.5671],\n",
      "        [0.4304, 0.6298, 0.5510],\n",
      "        [0.4671, 0.5910, 0.5266],\n",
      "        [0.4177, 0.6503, 0.5645]])\n"
     ]
    }
   ],
   "source": [
    "all_context_vecs = attn_weights @ inputs\n",
    "print(all_context_vecs)\n",
    "#重复了上一个操作"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25b245b8-7732-4fab-aa1c-e3d333195605",
   "metadata": {},
   "source": [
    "-  作为合理性检查，之前计算的上下文向量 $z^{(2)} = [0.4419, 0.6515, 0.5683]$ 可以在上图的第二行找到："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2570eb7d-aee1-457a-a61e-7544478219fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Previous 2nd context vector: tensor([0.4419, 0.6515, 0.5683])\n"
     ]
    }
   ],
   "source": [
    "print(\"Previous 2nd context vector:\", context_vec_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a303b6fb-9f7e-42bb-9fdb-2adabf0a6525",
   "metadata": {},
   "source": [
    "## 3.4 可调整参数的自注意力机制"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88363117-93d8-41fb-8240-f7cfe08b14a3",
   "metadata": {},
   "source": [
    "- 以下的概念框架展示了本节中开发的自注意力机制,以及这种机制是如何融入本书和本章的整体叙述与结构。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac9492ba-6f66-4f65-bd1d-87cf16d59928",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/13.webp\" width=\"400px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b90a77e-d746-4704-9354-1ddad86e6298",
   "metadata": {},
   "source": [
    "### 3.4.1 手把手的计算attention的值"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46e95a46-1f67-4b71-9e84-8e2db84ab036",
   "metadata": {},
   "source": [
    "- 在本节中，我们实现了原始 Transformer 架构、GPT 模型以及大多数流行 LLM 中使用的自注意力机制。  \n",
    "- 这种自注意力机制被称为“缩放点积注意力”（scaled dot-product attention）。  \n",
    "- 整体思路与之前相似：  \n",
    "  - 我们希望计算针对特定输入元素的上下文向量，即输入向量的加权和。  \n",
    "  - 为此，我们需要生成注意力权重。  \n",
    "- 如你所见，与之前介绍的基本注意力机制相比，只有一些细微差异：  \n",
    "  - 最显著的区别是引入了在模型训练过程中更新的权重矩阵。  \n",
    "  - 这些可训练的权重矩阵至关重要，它们使模型（尤其是注意力模块）能够学习生成“优质”的上下文向量。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59db4093-93e8-4bee-be8f-c8fac8a08cdd",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/14.webp\" width=\"600px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d996671-87aa-45c9-b2e0-07a7bcc9060a",
   "metadata": {},
   "source": [
    "- 按照步骤实现自注意力机制，我们将首先介绍三个训练权重矩阵 $W_q$、$W_k$ 和 $W_v$。  \n",
    "- 这三个矩阵用于通过矩阵乘法将嵌入的输入标记 $x^{(i)}$ 映射到查询向量、键向量和值向量：\n",
    "- (译者: 分别是Query、Key、Value,专有名词)   \n",
    "\n",
    "  - 查询向量：$q^{(i)} = W_q \\,x^{(i)}$  \n",
    "  - 键向量：$k^{(i)} = W_k \\,x^{(i)}$  \n",
    "  - 值向量：$v^{(i)} = W_v \\,x^{(i)}$  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f334313-5fd0-477b-8728-04080a427049",
   "metadata": {},
   "source": [
    "- 输入 $x$ 和查询向量 $q$ 的嵌入维度可以相同，也可以不同，具体取决于模型的设计和实现方式。\n",
    "- 在 GPT 模型中，输入和输出维度通常是相同的，但为了便于示范并更好地理解计算过程，这里我们选择了不同的输入和输出维度："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8250fdc6-6cd6-4c5b-b9c0-8c643aadb7db",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_2 = inputs[1] # second input element\n",
    "d_in = inputs.shape[1] # the input embedding size, d=3\n",
    "d_out = 2 # the output embedding size, d=2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f528cfb3-e226-47dd-b363-cc2caaeba4bf",
   "metadata": {},
   "source": [
    "- 下面，我们初始化三个权重矩阵；请注意，为了简化输出并便于示范，我们将 `requires_grad=False`\n",
    "- 但如果我们要在模型训练中使用这些权重矩阵，应将 `requires_grad=True`，以便在训练过程中更新这些矩阵。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "bfd7259a-f26c-4cea-b8fc-282b5cae1e00",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(123)\n",
    "#固定随机种子确保可复现性\n",
    "\n",
    "W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)\n",
    "W_key   = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)\n",
    "W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)\n",
    "#初始化三个矩阵来存放\n",
    "#不要求梯度降低了复杂度"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abfd0b50-7701-4adb-821c-e5433622d9c4",
   "metadata": {},
   "source": [
    "- 计算这三个向量值"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "73cedd62-01e1-4196-a575-baecc6095601",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.4306, 1.4551])\n"
     ]
    }
   ],
   "source": [
    "query_2 = x_2 @ W_query # _2 because it's with respect to the 2nd input element\n",
    "key_2 = x_2 @ W_key \n",
    "value_2 = x_2 @ W_value\n",
    "#点积计算\n",
    "print(query_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9be308b3-aca3-421b-b182-19c3a03b71c7",
   "metadata": {},
   "source": [
    "- 我们可以清晰地看到,embedding被降维了:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8c1c3949-fc08-4d19-a41e-1c235b4e631b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "keys.shape: torch.Size([6, 2])\n",
      "values.shape: torch.Size([6, 2])\n"
     ]
    }
   ],
   "source": [
    "keys = inputs @ W_key \n",
    "values = inputs @ W_value\n",
    "\n",
    "print(\"keys.shape:\", keys.shape)\n",
    "print(\"values.shape:\", values.shape)\n",
    "#中途检验下"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bac5dfd6-ade8-4e7b-b0c1-bed40aa24481",
   "metadata": {},
   "source": [
    "- 在下一步 **步骤 2** 中，我们通过计算查询向量和每个键向量之间的点积来计算未归一化的注意力得分："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ed0a2b7-5c50-4ede-90cf-7ad74412b3aa",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/15.webp\" width=\"600px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "64cbc253-a182-4490-a765-246979ea0a28",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1.8524)\n"
     ]
    }
   ],
   "source": [
    "keys_2 = keys[1] # Python starts index at 0\n",
    "attn_score_22 = query_2.dot(keys_2)\n",
    "print(attn_score_22)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e9d15c0-c24e-4e6f-a160-6349b418f935",
   "metadata": {},
   "source": [
    "- 因为我们有六个输入,所以我们有六个attention score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b14e44b5-d170-40f9-8847-8990804af26d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440])\n"
     ]
    }
   ],
   "source": [
    "attn_scores_2 = query_2 @ keys.T # All attention scores for given query\n",
    "print(attn_scores_2)\n",
    "#计算注意力跟query值"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8622cf39-155f-4eb5-a0c0-82a03ce9b999",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/16.webp\" width=\"600px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1609edb-f089-461a-8de2-c20c1bb29836",
   "metadata": {},
   "source": [
    "- 接下来，在 **步骤 3** 中，我们使用之前提到的 softmax 函数计算注意力权重（归一化后的注意力得分，总和为 1）。\n",
    "- 与之前的不同之处在于，我们现在通过将注意力得分除以嵌入维度的平方根 $\\sqrt{d_k}$（即 `d_k**0.5`）来对注意力得分进行缩放："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "146f5587-c845-4e30-9894-c7ed3a248153",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.1500, 0.2264, 0.2199, 0.1311, 0.0906, 0.1820])\n"
     ]
    }
   ],
   "source": [
    "d_k = keys.shape[1]\n",
    "attn_weights_2 = torch.softmax(attn_scores_2 / d_k**0.5, dim=-1)\n",
    "#压缩函数, 有利于储存与比较\n",
    "print(attn_weights_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8f61a28-b103-434a-aee1-ae7cbd821126",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/17.webp\" width=\"600px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1890e3f9-db86-4ab8-9f3b-53113504a61f",
   "metadata": {},
   "source": [
    "- 在**第四步**, 我们可以计算每一个token的向量了:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e138f033-fa7e-4e3a-8764-b53a96b26397",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.3061, 0.8210])\n"
     ]
    }
   ],
   "source": [
    "context_vec_2 = attn_weights_2 @ values\n",
    "print(context_vec_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d7b2907-e448-473e-b46c-77735a7281d8",
   "metadata": {},
   "source": [
    "### 3.4.2 自注意模块"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04313410-3155-4d90-a7a3-2f3386e73677",
   "metadata": {},
   "source": [
    "- 下面是代码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "51590326-cdbe-4e62-93b1-17df71c11ee4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2996, 0.8053],\n",
      "        [0.3061, 0.8210],\n",
      "        [0.3058, 0.8203],\n",
      "        [0.2948, 0.7939],\n",
      "        [0.2927, 0.7891],\n",
      "        [0.2990, 0.8040]], grad_fn=<MmBackward0>)\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class SelfAttention_v1(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out):\n",
    "        super().__init__()\n",
    "        self.W_query = nn.Parameter(torch.rand(d_in, d_out))\n",
    "        self.W_key   = nn.Parameter(torch.rand(d_in, d_out))\n",
    "        self.W_value = nn.Parameter(torch.rand(d_in, d_out))\n",
    "        #定义QKV的随机矩阵\n",
    "\n",
    "    def forward(self, x):\n",
    "        keys = x @ self.W_key\n",
    "        queries = x @ self.W_query\n",
    "        values = x @ self.W_value\n",
    "        \n",
    "        attn_scores = queries @ keys.T # omega\n",
    "        attn_weights = torch.softmax(\n",
    "            attn_scores / keys.shape[-1]**0.5, dim=-1\n",
    "        )\n",
    "        #模型的训练传递\n",
    "        context_vec = attn_weights @ values\n",
    "        return context_vec\n",
    "\n",
    "torch.manual_seed(123)\n",
    "sa_v1 = SelfAttention_v1(d_in, d_out)\n",
    "print(sa_v1(inputs))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ee1a024-84a5-425a-9567-54ab4e4ed445",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/18.webp\" width=\"400px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "048e0c16-d911-4ec8-b0bc-45ceec75c081",
   "metadata": {},
   "source": [
    "- 我们可以使用 PyTorch 的 `Linear` 层简化上述实现，禁用偏置项后，`Linear` 层相当于矩阵乘法。\n",
    "- 使用 `nn.Linear` 替代手动使用 `nn.Parameter(torch.rand(...))` 的一个主要优势是，`nn.Linear` 具有推荐的权重初始化方案，这有助于模型训练更加稳定。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "73f411e3-e231-464a-89fe-0a9035e5f839",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.0739,  0.0713],\n",
      "        [-0.0748,  0.0703],\n",
      "        [-0.0749,  0.0702],\n",
      "        [-0.0760,  0.0685],\n",
      "        [-0.0763,  0.0679],\n",
      "        [-0.0754,  0.0693]], grad_fn=<MmBackward0>)\n"
     ]
    }
   ],
   "source": [
    "class SelfAttention_v2(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        #权重初始化\n",
    "\n",
    "    def forward(self, x):\n",
    "        keys = self.W_key(x)\n",
    "        queries = self.W_query(x)\n",
    "        values = self.W_value(x)\n",
    "        \n",
    "        attn_scores = queries @ keys.T \n",
    "        #Query跟Key的计算 得出初始的分数传递到后面进行归一化操作\n",
    "        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)\n",
    "\n",
    "        context_vec = attn_weights @ values\n",
    "        #直接基于注意力对于文本计算\n",
    "        return context_vec\n",
    "\n",
    "torch.manual_seed(789)\n",
    "sa_v2 = SelfAttention_v2(d_in, d_out)\n",
    "print(sa_v2(inputs))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "915cd8a5-a895-42c9-8b8e-06b5ae19ffce",
   "metadata": {},
   "source": [
    "- `SelfAttention_v1` 和 `SelfAttention_v2` 会给出不同的输出，因为它们使用了不同的初始权重矩阵。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5025b37-0f2c-4a67-a7cb-1286af7026ab",
   "metadata": {},
   "source": [
    "## 3.5 对未出现的信息的隐藏"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aef0a6b8-205a-45bf-9d26-8fd77a8a03c3",
   "metadata": {},
   "source": [
    "- 在casual attention，对角线以上的注意力权重被掩蔽，确保在计算上下文向量时，LLM 无法利用位置的信息来调整注意力权重。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71e91bb5-5aae-4f05-8a95-973b3f988a35",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/19.webp\" width=\"400px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82f405de-cd86-4e72-8f3c-9ea0354946ba",
   "metadata": {},
   "source": [
    "### 3.5.1 因果自注意力机制"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "014f28d0-8218-48e4-8b9c-bdc5ce489218",
   "metadata": {},
   "source": [
    "- 在这一节中，我们将把之前的自注意力机制转换为因果自注意力机制。\n",
    "- 因果自注意力确保模型在预测序列中某个位置的值时，仅依赖于前面已知位置的输出，而不依赖于后续位置。\n",
    "- 换句话说，这确保了每个下一个词的预测仅依赖于前面的词。\n",
    "- 为了实现这一点，对于每个给定的标记，我们会将“未知的信息”（即输入文本中当前token之后的token）掩蔽掉："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57f99af3-32bc-48f5-8eb4-63504670ca0a",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/20.webp\" width=\"600px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbfaec7a-68f2-4157-a4b5-2aeceed199d9",
   "metadata": {},
   "source": [
    "- 为了说明和实现因果自注意力，让我们使用上一节中的注意力得分和权重："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "1933940d-0fa5-4b17-a3ce-388e5314a1bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.1921, 0.1646, 0.1652, 0.1550, 0.1721, 0.1510],\n",
      "        [0.2041, 0.1659, 0.1662, 0.1496, 0.1665, 0.1477],\n",
      "        [0.2036, 0.1659, 0.1662, 0.1498, 0.1664, 0.1480],\n",
      "        [0.1869, 0.1667, 0.1668, 0.1571, 0.1661, 0.1564],\n",
      "        [0.1830, 0.1669, 0.1670, 0.1588, 0.1658, 0.1585],\n",
      "        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],\n",
      "       grad_fn=<SoftmaxBackward0>)\n"
     ]
    }
   ],
   "source": [
    "# Reuse the query and key weight matrices of the\n",
    "# SelfAttention_v2 object from the previous section for convenience\n",
    "queries = sa_v2.W_query(inputs)\n",
    "keys = sa_v2.W_key(inputs) \n",
    "attn_scores = queries @ keys.T\n",
    "\n",
    "attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)\n",
    "print(attn_weights)\n",
    "#用mask的数据重新算了一次"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89020a96-b34d-41f8-9349-98c3e23fd5d6",
   "metadata": {},
   "source": [
    "- 隐藏未知信息的attention score最简单的方法是通过 PyTorch 的 `tril` 函数进行掩蔽，其中主对角线以下的元素（包括对角线本身）设置为 1，主对角线以上的元素设置为 0："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "43f3d2e3-185b-4184-9f98-edde5e6df746",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 0., 0., 0., 0., 0.],\n",
      "        [1., 1., 0., 0., 0., 0.],\n",
      "        [1., 1., 1., 0., 0., 0.],\n",
      "        [1., 1., 1., 1., 0., 0.],\n",
      "        [1., 1., 1., 1., 1., 0.],\n",
      "        [1., 1., 1., 1., 1., 1.]])\n"
     ]
    }
   ],
   "source": [
    "context_length = attn_scores.shape[0]\n",
    "mask_simple = torch.tril(torch.ones(context_length, context_length))\n",
    "#Mask矩阵,直接保留Diagonal下部分的,上部分掩盖掉\n",
    "print(mask_simple)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efce2b08-3583-44da-b3fc-cabdd38761f6",
   "metadata": {},
   "source": [
    "- 然后，我们可以将注意力权重与这个mask相乘，以将对角线以上的注意力得分置为零："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "9f531e2e-f4d2-4fea-a87f-4c132e48b9e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.1921, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.2041, 0.1659, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.2036, 0.1659, 0.1662, 0.0000, 0.0000, 0.0000],\n",
      "        [0.1869, 0.1667, 0.1668, 0.1571, 0.0000, 0.0000],\n",
      "        [0.1830, 0.1669, 0.1670, 0.1588, 0.1658, 0.0000],\n",
      "        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],\n",
      "       grad_fn=<MulBackward0>)\n"
     ]
    }
   ],
   "source": [
    "masked_simple = attn_weights*mask_simple\n",
    "print(masked_simple)\n",
    "#简单的效果图"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3eb35787-cf12-4024-b66d-e7215e175500",
   "metadata": {},
   "source": [
    "- 如果在 softmax 之后进行掩蔽，它会破坏 softmax 所创建的概率分布。\n",
    "- softmax 确保所有输出值的总和为 1。\n",
    "- 如果在 softmax 之后进行掩蔽，就需要重新归一化输出，确保其总和为 1，这会使过程更加复杂，并可能带来意想不到的效果。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94db92d7-c397-4e42-bd8a-6a2b3e237e0f",
   "metadata": {},
   "source": [
    "- 我们可以用以下方式确保所有的数据都是归一化的"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "6d392083-fd81-4f70-9bdf-8db985e673d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],\n",
      "        [0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],\n",
      "        [0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],\n",
      "        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],\n",
      "       grad_fn=<DivBackward0>)\n"
     ]
    }
   ],
   "source": [
    "row_sums = masked_simple.sum(dim=-1, keepdim=True)\n",
    "masked_simple_norm = masked_simple / row_sums\n",
    "print(masked_simple_norm)\n",
    "#掩码之后的softmax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "512e7cf4-dc0e-4cec-948e-c7a3c4eb6877",
   "metadata": {},
   "source": [
    "- 尽管我们在技术上已经完成了因果注意力机制的编码，但让我们简要地探讨一种更高效的方法，以实现与上述相同的效果。\n",
    "- 因此，在注意力得分进入 softmax 函数之前，我们可以将对角线以上的未归一化注意力得分用负无穷大进行掩蔽，而不是将其置零并重新归一化："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb682900-8df2-4767-946c-a82bee260188",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/21.webp\" width=\"450px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "a2be2f43-9cf0-44f6-8d8b-68ef2fb3cc39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2899,   -inf,   -inf,   -inf,   -inf,   -inf],\n",
      "        [0.4656, 0.1723,   -inf,   -inf,   -inf,   -inf],\n",
      "        [0.4594, 0.1703, 0.1731,   -inf,   -inf,   -inf],\n",
      "        [0.2642, 0.1024, 0.1036, 0.0186,   -inf,   -inf],\n",
      "        [0.2183, 0.0874, 0.0882, 0.0177, 0.0786,   -inf],\n",
      "        [0.3408, 0.1270, 0.1290, 0.0198, 0.1290, 0.0078]],\n",
      "       grad_fn=<MaskedFillBackward0>)\n"
     ]
    }
   ],
   "source": [
    "mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)\n",
    "#创建一个全1的三角,去上部分变成0\n",
    "masked = attn_scores.masked_fill(mask.bool(), -torch.inf)\n",
    "#有掩码的地方变为负无穷\n",
    "print(masked)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91d5f803-d735-4543-b9da-00ac10fb9c50",
   "metadata": {},
   "source": [
    "- 结果显然是归一化的"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "b1cd6d7f-16f2-43c1-915e-0824f1a4bc52",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],\n",
      "        [0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],\n",
      "        [0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],\n",
      "        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],\n",
      "       grad_fn=<SoftmaxBackward0>)\n"
     ]
    }
   ],
   "source": [
    "attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=-1)\n",
    "print(attn_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7636fc5f-6bc6-461e-ac6a-99ec8e3c0912",
   "metadata": {},
   "source": [
    "### 3.5.2 使用Dropout防止过拟合"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec3dc7ee-6539-4fab-804a-8f31a890c85a",
   "metadata": {},
   "source": [
    "- 此外，我们还应用了丢弃（Dropout）来减少训练过程中的过拟合。\n",
    "- dropout可以应用于多个位置：\n",
    "  - 例如，在计算注意力权重之后；\n",
    "  - 或在将注意力权重与值向量相乘之后。\n",
    "- 在这里，我们选择在计算注意力权重之后应用丢弃掩码，因为这种做法更为常见。\n",
    "\n",
    "- 另外，在此示例中，我们使用了50%的丢弃率，这意味着随机屏蔽掉一半的注意力权重。（在后续训练GPT模型时，我们会使用更低的丢弃率，例如0.1或0.2。）"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee799cf6-6175-45f2-827e-c174afedb722",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/22.webp\" width=\"400px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a575458-a6da-4e54-8688-83e155f2de06",
   "metadata": {},
   "source": [
    "- 如果我们应用0.5（50%）的丢弃率，未被抛弃的值将相应地被缩放一个因子，1/0.5 = 2。\n",
    "- 这种缩放通过公式 1 / (1 - `dropout_rate`) 计算得出。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "0de578db-8289-41d6-b377-ef645751e33f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[2., 2., 2., 2., 2., 2.],\n",
      "        [0., 2., 0., 0., 0., 0.],\n",
      "        [0., 0., 2., 0., 2., 0.],\n",
      "        [2., 2., 0., 0., 0., 2.],\n",
      "        [2., 0., 0., 0., 0., 2.],\n",
      "        [0., 2., 0., 0., 0., 0.]])\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(123)\n",
    "dropout = torch.nn.Dropout(0.5) \n",
    "# dropout rate of 50%丢包率doge\n",
    "example = torch.ones(6, 6) \n",
    "# create a matrix of ones满的6*6矩阵被1包圆了\n",
    "\n",
    "print(dropout(example))\n",
    "#输出需要被放大相应的倍数,为了维持恒定"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "b16c5edb-942b-458c-8e95-25e4e355381e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[2.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.8966, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.0000, 0.6206, 0.0000, 0.0000, 0.0000],\n",
      "        [0.5517, 0.4921, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.4350, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.3327, 0.0000, 0.0000, 0.0000, 0.0000]],\n",
      "       grad_fn=<MulBackward0>)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(123)\n",
    "print(dropout(attn_weights))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "269df5c8-3e25-49d0-95d3-bb232287404f",
   "metadata": {},
   "source": [
    "- 生成的输出可能会因操作系统的不同而有所不同；\n",
    "- 你可以在 [PyTorch 问题追踪器](https://github.com/pytorch/pytorch/issues/121595) 上了解更多内容。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdc14639-5f0f-4840-aa9d-8eb36ea90fb7",
   "metadata": {},
   "source": [
    "### 3.5.3 实现一个简洁的因果自注意力"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09c41d29-1933-43dc-ada6-2dbb56287204",
   "metadata": {},
   "source": [
    "- 现在，我们准备实现一个完整的自注意力机制，包含因果掩码和dropout。\n",
    "- 另一项任务是实现代码以处理包含多个输入的批次，确保我们的 `CausalAttention` 类能够支持第二章中实现的数据加载器所生成的批量输出。\n",
    "- 为了简化起见，我们通过复制输入文本示例来模拟批量输入："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "977a5fa7-a9d5-4e2e-8a32-8e0331ccfe28",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 6, 3])\n"
     ]
    }
   ],
   "source": [
    "batch = torch.stack((inputs, inputs), dim=0)\n",
    "#相同的tensor按照指定维度堆叠\n",
    "print(batch.shape) # 2 inputs with 6 tokens each, and each token has embedding dimension 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0de8253-5387-4dd1-b53b-72f8f10b9ecb",
   "metadata": {},
   "source": [
    "- 缩放因子  \\sqrt{d}  的引入解决了注意力机制中的数值不稳定问题。\n",
    "- 它确保了即使嵌入维度  d  较大，点积得分也能被合理地控制在一个适当范围，方便 Softmax 生成平滑的注意力分布，且梯度不会过大或过小。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "60d8c2eb-2d8e-4d2c-99bc-9eef8cc53ca0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[-0.4519,  0.2216],\n",
      "         [-0.5874,  0.0058],\n",
      "         [-0.6300, -0.0632],\n",
      "         [-0.5675, -0.0843],\n",
      "         [-0.5526, -0.0981],\n",
      "         [-0.5299, -0.1081]],\n",
      "\n",
      "        [[-0.4519,  0.2216],\n",
      "         [-0.5874,  0.0058],\n",
      "         [-0.6300, -0.0632],\n",
      "         [-0.5675, -0.0843],\n",
      "         [-0.5526, -0.0981],\n",
      "         [-0.5299, -0.1081]]], grad_fn=<UnsafeViewBackward0>)\n",
      "context_vecs.shape: torch.Size([2, 6, 2])\n"
     ]
    }
   ],
   "source": [
    "class CausalAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, context_length,\n",
    "                 dropout, qkv_bias=False):\n",
    "        #初始化定义网络结构和参数\n",
    "        super().__init__()\n",
    "        self.d_out = d_out\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.dropout = nn.Dropout(dropout) # New\n",
    "        self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1)) # New\n",
    "        #定义QKV并对进行dropout防止过拟合\n",
    "        #注册mask向量,对未来进行负无穷的拟合\n",
    "    def forward(self, x):\n",
    "        b, num_tokens, d_in = x.shape # New batch dimension b\n",
    "        #提取batch的大小、token的数量、跟宽度\n",
    "        keys = self.W_key(x)\n",
    "        queries = self.W_query(x)\n",
    "        values = self.W_value(x)\n",
    "        #进行运算计算\n",
    "        attn_scores = queries @ keys.transpose(1, 2) # Changed transpose\n",
    "        #通过点积来计算attention的数值\n",
    "        attn_scores.masked_fill_(  # New, _ ops are in-place\n",
    "            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)  # `:num_tokens` to account for cases where the number of tokens in the batch is smaller than the supported context_size\n",
    "        attn_weights = torch.softmax(\n",
    "            attn_scores / keys.shape[-1]**0.5, dim=-1## 缩放因子 √d，用于稳定梯度\n",
    "        )\n",
    "        #在时间顺序上进行mask确保信息不会被泄露\n",
    "        attn_weights = self.dropout(attn_weights) # New\n",
    "        #防止过拟合的dropout处理方式\n",
    "        context_vec = attn_weights @ values\n",
    "        # 根据注意力权重计算上下文向量\n",
    "        return context_vec\n",
    "\n",
    "torch.manual_seed(123)\n",
    "context_length = batch.shape[1]\n",
    "ca = CausalAttention(d_in, d_out, context_length, 0.0)\n",
    "\n",
    "context_vecs = ca(batch)\n",
    "\n",
    "print(context_vecs)\n",
    "print(\"context_vecs.shape:\", context_vecs.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4333d12-17e4-4bb5-9d83-54b3a32618cd",
   "metadata": {},
   "source": [
    "- Dropout仅在训练时要使用,验证时不需要"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a554cf47-558c-4f45-84cd-bf9b839a8d50",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/23.webp\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8bef90f-cfd4-4289-b0e8-6a00dc9be44c",
   "metadata": {},
   "source": [
    "## 3.6 拓展单头至多方注意"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11697757-9198-4a1c-9cee-f450d8bbd3b9",
   "metadata": {},
   "source": [
    "### 3.6.1 堆叠多个单头注意力层"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70766faf-cd53-41d9-8a17-f1b229756a5a",
   "metadata": {},
   "source": [
    "- 以下是之前实现的自注意力机制总结（为简化起见，未展示因果和dropout掩码）：\n",
    "\n",
    "- 这种机制也称为单头注意力：\n",
    "\n",
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/24.webp\" width=\"400px\">\n",
    "\n",
    "- 我们通过堆叠多个单头注意力模块来构建多头注意力模块：\n",
    "\n",
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/25.webp\" width=\"400px\">\n",
    "\n",
    "- 多头注意力的核心思想是使用不同的学习到的线性投影，并行地多次运行注意力机制。这使得模型能够在不同位置同时关注来自不同表示子空间的信息。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77a341bf-122a-4e31-99db-39259cf8d1b2",
   "metadata": {},
   "source": [
    "在 Python 中，super().__init__() 是一种调用父类（基类）构造函数的方法，常用于类继承的场景中。它确保子类能够正确初始化父类的属性和方法。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "b9a66e11-7105-4bb4-be84-041f1a1f3bd2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[-0.4519,  0.2216,  0.4772,  0.1063],\n",
      "         [-0.5874,  0.0058,  0.5891,  0.3257],\n",
      "         [-0.6300, -0.0632,  0.6202,  0.3860],\n",
      "         [-0.5675, -0.0843,  0.5478,  0.3589],\n",
      "         [-0.5526, -0.0981,  0.5321,  0.3428],\n",
      "         [-0.5299, -0.1081,  0.5077,  0.3493]],\n",
      "\n",
      "        [[-0.4519,  0.2216,  0.4772,  0.1063],\n",
      "         [-0.5874,  0.0058,  0.5891,  0.3257],\n",
      "         [-0.6300, -0.0632,  0.6202,  0.3860],\n",
      "         [-0.5675, -0.0843,  0.5478,  0.3589],\n",
      "         [-0.5526, -0.0981,  0.5321,  0.3428],\n",
      "         [-0.5299, -0.1081,  0.5077,  0.3493]]], grad_fn=<CatBackward0>)\n",
      "context_vecs.shape: torch.Size([2, 6, 4])\n"
     ]
    }
   ],
   "source": [
    "class MultiHeadAttentionWrapper(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):\n",
    "        super().__init__() \n",
    "        #多个实例,每个都是一个头\n",
    "        self.heads = nn.ModuleList(\n",
    "            [CausalAttention(d_in, d_out, context_length, dropout, qkv_bias) \n",
    "             for _ in range(num_heads)]\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.cat([head(x) for head in self.heads], dim=-1)\n",
    "    #模型的训练\n",
    "\n",
    "\n",
    "torch.manual_seed(123)\n",
    "\n",
    "context_length = batch.shape[1] # This is the number of tokens\n",
    "d_in, d_out = 3, 2\n",
    "mha = MultiHeadAttentionWrapper(\n",
    "    d_in, d_out, context_length, 0.0, num_heads=2\n",
    ")\n",
    "\n",
    "context_vecs = mha(batch)\n",
    "\n",
    "print(context_vecs)\n",
    "print(\"context_vecs.shape:\", context_vecs.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "193d3d2b-2578-40ba-b791-ea2d49328e48",
   "metadata": {},
   "source": [
    "- 在上述实现中，嵌入维度为4，因为我们将 `d_out=2` 作为键、查询和值向量以及上下文向量的嵌入维度。由于使用了2个注意力头，输出嵌入维度为 2 * 2 = 4。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6836b5da-ef82-4b4c-bda1-72a462e48d4e",
   "metadata": {},
   "source": [
    "### 3.6.2 利用权重拆分实现多头注意力"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4b48d0d-71ba-4fa0-b714-ca80cabcb6f7",
   "metadata": {},
   "source": [
    "- 尽管上述实现是一种直观且功能完整的多头注意力机制（通过封装之前的单头注意力 `CausalAttention` 实现），我们仍可以编写一个独立的 `MultiHeadAttention` 类来实现相同的功能。\n",
    "\n",
    "- 在这个独立的 `MultiHeadAttention` 类中，我们不会将单个注意力头进行拼接。\n",
    "- 相反，我们会创建独立的 W_query、W_key 和 W_value 权重矩阵，并将它们拆分为每个注意力头的单独矩阵："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "110b0188-6e9e-4e56-a988-10523c6c8538",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0.3190, 0.4858],\n",
      "         [0.2943, 0.3897],\n",
      "         [0.2856, 0.3593],\n",
      "         [0.2693, 0.3873],\n",
      "         [0.2639, 0.3928],\n",
      "         [0.2575, 0.4028]],\n",
      "\n",
      "        [[0.3190, 0.4858],\n",
      "         [0.2943, 0.3897],\n",
      "         [0.2856, 0.3593],\n",
      "         [0.2693, 0.3873],\n",
      "         [0.2639, 0.3928],\n",
      "         [0.2575, 0.4028]]], grad_fn=<ViewBackward0>)\n",
      "context_vecs.shape: torch.Size([2, 6, 2])\n"
     ]
    }
   ],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        assert (d_out % num_heads == 0), \\\n",
    "        \"d_out must be divisible by num_heads\"\n",
    "        #确保是可以被整除的\n",
    "            \n",
    "\n",
    "        self.d_out = d_out\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim\n",
    "        #初始化头的维度、数量\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs\n",
    "        #头的输出结合线性层\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        #进行dropout防止过拟合\n",
    "        self.register_buffer(\n",
    "            \"mask\",\n",
    "            torch.triu(torch.ones(context_length, context_length),\n",
    "                       diagonal=1)\n",
    "        )\n",
    "        # 上三角掩码，确保因果性\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, num_tokens, d_in = x.shape\n",
    "\n",
    "        keys = self.W_key(x) # Shape: (b, num_tokens, d_out)\n",
    "        queries = self.W_query(x)\n",
    "        values = self.W_value(x)\n",
    "        #把输出的维度拆成头*头大小\n",
    "        # We implicitly split the matrix by adding a `num_heads` dimension\n",
    "        # Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)\n",
    "        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) \n",
    "        values = values.view(b, num_tokens, self.num_heads, self.head_dim)\n",
    "        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)\n",
    "        #转制维度,听说是为了更好的计算注意力\n",
    "        # Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)\n",
    "        keys = keys.transpose(1, 2)\n",
    "        queries = queries.transpose(1, 2)\n",
    "        values = values.transpose(1, 2)\n",
    "        # 计算缩放点积注意力\n",
    "        # Compute scaled dot-product attention (aka self-attention) with a causal mask\n",
    "        attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head\n",
    "        # 将掩码缩减到当前 token 数量，并转换为布尔型\n",
    "        # 进而实现动态遮蔽,所以不用另开好几个数组\n",
    "        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]\n",
    "        # 遮蔽矩阵\n",
    "        # Use the mask to fill attention scores\n",
    "        attn_scores.masked_fill_(mask_bool, -torch.inf)\n",
    "        #归一化\n",
    "        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)\n",
    "        attn_weights = self.dropout(attn_weights)\n",
    "\n",
    "        # Shape: (b, num_tokens, num_heads, head_dim)\n",
    "        context_vec = (attn_weights @ values).transpose(1, 2) \n",
    "        #头的合并\n",
    "        # Combine heads, where self.d_out = self.num_heads * self.head_dim\n",
    "        #对上下文向量的形状进行调整，确保输出的形状\n",
    "        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)\n",
    "        context_vec = self.out_proj(context_vec) # optional projection\n",
    "\n",
    "        return context_vec\n",
    "\n",
    "torch.manual_seed(123)\n",
    "\n",
    "batch_size, context_length, d_in = batch.shape\n",
    "d_out = 2\n",
    "mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2)\n",
    "\n",
    "context_vecs = mha(batch)\n",
    "\n",
    "print(context_vecs)\n",
    "print(\"context_vecs.shape:\", context_vecs.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d334dfb5-2b6c-4c33-82d5-b4e9db5867bb",
   "metadata": {},
   "source": [
    "- 请注意，上述实现本质上是 `MultiHeadAttentionWrapper` 的重写版本，并且更加高效。  \n",
    "- 生成的输出看起来略有不同，因为随机权重初始化有所不同，但两者都是完全可用的实现，可以在我们将在后续章节中实现的 GPT 类中使用。  \n",
    "- 另外，值得注意的是，我们在上面的 `MultiHeadAttention` 类中添加了一个线性投影层（`self.out_proj`）。这只是一个线性变换，不改变维度。在 LLM 实现中，使用这样的投影层是标准做法，但它并非严格必要（近期的研究表明，去除该层不会影响模型的表现；请参阅本章末尾的进一步阅读部分）。  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbe5d396-c990-45dc-9908-2c621461f851",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/ch03_compressed/26.webp\" width=\"400px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b0ed78c-e8ac-4f8f-a479-a98242ae8f65",
   "metadata": {},
   "source": [
    "- 请注意，如果你对上述内容的紧凑和高效实现感兴趣，可以考虑使用PyTorch中的 [`torch.nn.MultiheadAttention`](https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html) 类。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "363701ad-2022-46c8-9972-390d2a2b9911",
   "metadata": {},
   "source": [
    "- 由于上述实现初看起来可能有些复杂，我们来看一下执行 `attn_scores = queries @ keys.transpose(2, 3)` 时会发生什么："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "e8cfc1ae-78ab-4faa-bc73-98bd054806c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[1.3208, 1.1631, 1.2879],\n",
      "          [1.1631, 2.2150, 1.8424],\n",
      "          [1.2879, 1.8424, 2.0402]],\n",
      "\n",
      "         [[0.4391, 0.7003, 0.5903],\n",
      "          [0.7003, 1.3737, 1.0620],\n",
      "          [0.5903, 1.0620, 0.9912]]]])\n"
     ]
    }
   ],
   "source": [
    "# (b, num_heads, num_tokens, head_dim) = (1, 2, 3, 4)\n",
    "a = torch.tensor([[[[0.2745, 0.6584, 0.2775, 0.8573],\n",
    "                    [0.8993, 0.0390, 0.9268, 0.7388],\n",
    "                    [0.7179, 0.7058, 0.9156, 0.4340]],\n",
    "\n",
    "                   [[0.0772, 0.3565, 0.1479, 0.5331],\n",
    "                    [0.4066, 0.2318, 0.4545, 0.9737],\n",
    "                    [0.4606, 0.5159, 0.4220, 0.5786]]]])\n",
    "\n",
    "print(a @ a.transpose(2, 3))\n",
    "#每个注意力头中，输出矩阵的值表示每个 token 对其他 token 的相关性。\n",
    "#模型可以计算 token 之间的相关性，为注意力分布的生成奠定基础"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0587b946-c8f2-4888-adbf-5a5032fbfd7b",
   "metadata": {},
   "source": [
    "- 在这种情况下，PyTorch 中的矩阵乘法实现将处理四维输入张量，使得矩阵乘法在最后两个维度（`num_tokens`, `head_dim`）之间进行，然后对各个头进行重复计算。  \n",
    "\n",
    "- 例如，以下方法提供了一种更紧凑的方式来分别计算每个头的矩阵乘法：  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "053760f1-1a02-42f0-b3bf-3d939e407039",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First head:\n",
      " tensor([[1.3208, 1.1631, 1.2879],\n",
      "        [1.1631, 2.2150, 1.8424],\n",
      "        [1.2879, 1.8424, 2.0402]])\n",
      "\n",
      "Second head:\n",
      " tensor([[0.4391, 0.7003, 0.5903],\n",
      "        [0.7003, 1.3737, 1.0620],\n",
      "        [0.5903, 1.0620, 0.9912]])\n"
     ]
    }
   ],
   "source": [
    "first_head = a[0, 0, :, :]\n",
    "#定义第一个头\n",
    "first_res = first_head @ first_head.T\n",
    "#第一个矩阵的自相关性\n",
    "print(\"First head:\\n\", first_res)\n",
    "\n",
    "second_head = a[0, 1, :, :]\n",
    "second_res = second_head @ second_head.T\n",
    "print(\"\\nSecond head:\\n\", second_res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dec671bf-7938-4304-ad1e-75d9920e7f43",
   "metadata": {},
   "source": [
    "# 总结与收获"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa3e4113-ffca-432c-b3ec-7a50bd15da25",
   "metadata": {},
   "source": [
    "- 请参阅 [./multihead-attention.ipynb](./multihead-attention.ipynb) 代码笔记本，它是数据加载器（第2章）的简洁版本，加上我们在本章实现的多头注意力类，后续章节中训练GPT模型时将需要使用。\n",
    "- 你可以在 [./exercise-solutions.ipynb](./exercise-solutions.ipynb) 中找到习题解答。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
